import cv2

import numpy as np
from PIL import Image
from AI_ICCN.yolo import YOLO

class ICCN_Det():
    def __init__(self):
        super(ICCN_Det).__init__()
        self.yolo = YOLO()
        self.class_names = ['FOREIGN', 'SOLDER', 'BRIDGE', 'XIZHU', 'SHIFT']
        self.OK_result = "OK"

    def start_det(self,datalist,ng_dict):
        res_dict = {}
        for data in datalist:
            key,image = data[0],data[2]
            image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
            res= self.yolo.detect_image(image)
            if res != 'OK' :
                ng_dict[key] = res
        return ng_dict

    def det_image(self,image):
        image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        image,ress  = self.yolo.detect_image_one(image)
        return image,ress